import gymnasium as gym
import cv2 as cv
import tensorflow as tf
from tensorflow.keras.layers import Input, Dense

class ActorNetwork(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.actor_mean = tf.keras.Sequential([
            Input(shape=(s_dim,)),
            Dense(64, activation=activation, kernel_initializer=initializer),
            Dense(64, activation=activation, kernel_initializer=initializer),
            Dense(a_dim, kernel_initializer=initializer)
        ])

    def call(self, state):
        mean = self.actor_mean(state)
        return mean

activation = tf.keras.activations.tanh
initializer = tf.keras.initializers.Orthogonal  # 가중치 초기화

task_name = 'Ant-v4'
env = gym.make(task_name, render_mode='rgb_array')
s_dim = env.observation_space.shape[0]  # 상태 공간
a_dim = env.action_space.shape[0]

actor = ActorNetwork()
actor.load_weights('f10-7.weights.h5')

length = 0
s, info = env.reset()
while True:
    mean = actor(s.reshape(1, -1))
    s, r, terminated, truncated, info = env.step(mean[0])  # 평균을 행동으로 취함(탐욕 선택)
    length += 1

    cv.imshow(task_name + ' animation',
              cv.cvtColor(env.render(), cv.COLOR_BGR2RGB))
    key = cv.waitKey(10)

    if terminated or truncated:
        print("에피소드의 길이:", length)
        break

env.close()
if cv.waitKey() == ord('q'):
    cv.destroyAllWindows()
